from pathology.controllers.mplib_planner import MPLibPlanner
import numpy as np
from scipy.spatial.transform import Rotation as R
import sapien


class SlidePlanner(MPLibPlanner):

    def __init__(self, scene:sapien.Scene, urdf, srdf=None, move_group='tool0',
                    pose_robotiq_to_tool0:sapien.Pose=None,
                    pose_camera_to_tool0:sapien.Pose=None):
        super().__init__(scene, urdf, srdf, move_group)
        self._scene = scene
        self.pose_slide_to_robotiq = sapien.Pose(p=[0.29, 0., 0.], q=[1, 0, 0, 0])
        self.pose_robotiq_to_tool0 = pose_robotiq_to_tool0
        self.pose_slide_to_tool0 = self.pose_robotiq_to_tool0 * self.pose_slide_to_robotiq
        self.pose_camera_to_tool0 = pose_camera_to_tool0

        self.phase4_waypoints = [np.array([-0.025, 0, -0.035, 0, 0, 0, 1]),
                                np.array([-0.02, 0, -0.035, 0, 0, 0, 1]),
                                np.array([-0.015, 0, -0.035, 0, 0, 0, 1]),
                                np.array([-0.01, 0., -0.03, 0, 0, 0, 1]),
                                np.array([-0.005, 0., -0.025, 0, 0, 0, 1]),
                                np.array([0.0, 0., -0.015, 0, 0, 0, 1]),
                                np.array([0.005, 0., -0.005, 0, 0, 0, 1]),
                                np.array([0.008, 0., 0.005, 0, 0, 0, 1]),
                            ]

    def plan_slide_pose(self, slide_target_pose, robot_qpos):
        tool0_target_pose = slide_target_pose * self.pose_slide_to_tool0.inv()
        result = self.planner.plan_pose(tool0_target_pose,
                                        robot_qpos,
                                        time_step=self._scene.get_timestep(),
                                    )
        return result
    
    def plan_camera_pose(self, camera_target_pose, robot_qpos):
        tool0_target_pose = camera_target_pose * self.pose_camera_to_tool0.inv()
        result = self.planner.plan_pose(tool0_target_pose,
                                        robot_qpos,
                                        time_step=self._scene.get_timestep(),
                                    )
        if result['status'] != "Success":
            print('--')
        return result
    
    def plan_slide_traj(self, slide_target_traj, robot_qpos):
        result = {}
        result['status'] = 'Success'
        result['position'] = []
        result['velocity'] = []
        current_qpos = robot_qpos
        for current_target_pose in slide_target_traj:
            current_result = self.plan_slide_pose(current_target_pose, current_qpos)
            if current_result['status'] != 'Success':
                result['status'] = 'Failed'
                return result
            else:
                result['position'].append(current_result['position'])
                result['velocity'].append(current_result['velocity'])
                current_qpos = current_result['position'][-1]
        
        result['position'] = np.concatenate(result['position'], axis=0) if len(result['position']) > 0 else []
        result['velocity'] = np.concatenate(result['velocity'], axis=0) if len(result['velocity']) > 0 else []
        return result

    def phase1(self, obs):
        # move the slide above the cut to start a traj
        cut_pose:sapien.Pose = obs['cut_pose']
        phase1_related_pose = sapien.Pose(p=[-0.1, 0., 0.1], q=R.from_euler('xyz', [0, 1.5708, 0]).as_quat()[[3, 0, 1, 2]])
        phase1_target_pose = cut_pose * phase1_related_pose

        result = self.plan_slide_pose(phase1_target_pose, obs['qpos'])
        return result
    
    def phase2(self, obs):
        # move the slide get close to the cut
        cut_pose:sapien.Pose = obs['cut_pose']
        phase1_related_pose = sapien.Pose(p=[-0.03, 0., 0.02], q=R.from_euler('xyz', [0, 1.5708, 0]).as_quat()[[3, 0, 1, 2]])
        phase1_target_pose = cut_pose * phase1_related_pose

        result = self.plan_slide_pose(phase1_target_pose, obs['qpos'])
        return result
    
    def phase3(self, obs):
        # put the slide down into the water
        cut_pose:sapien.Pose = obs['cut_pose']
        phase1_related_pose = sapien.Pose(p=[-0.025, 0., -0.035], q=R.from_euler('xyz', [0, 1.5708, 0]).as_quat()[[3, 0, 1, 2]])
        phase1_target_pose = cut_pose * phase1_related_pose

        result = self.plan_slide_pose(phase1_target_pose, obs['qpos'])
        return result
    
    def phase4(self, obs):
        head_down_pose = sapien.Pose(p=[0., 0., 0.], q=R.from_euler('xyz', [0, 1.5708, 0]).as_quat()[[3, 0, 1, 2]])
        cut_pose:sapien.Pose = obs['cut_pose']
        traj_waypoints = []
        for waypoint in self.phase4_waypoints:
            waypoint_pose = sapien.Pose(p=waypoint[:3], q=waypoint[[-1, 3, 4, 5]])
            traj_waypoints.append(cut_pose * waypoint_pose * head_down_pose) 
        
        result = self.plan_slide_traj(traj_waypoints, obs['qpos'])
        return result

    def target_pose_from_2d(self, x, y, theta, h):
        # 从2D的x, y, theta和高度h计算出3D的目标位姿
        # 夹爪垂直朝下
        # 暂时没什么用
        target_pose = sapien.Pose(p=[x, y, h], q=R.from_euler('xyz', [0, np.pi/2, theta]).as_quat()[[3, 0, 1, 2]])
        return target_pose
    
    def phase0_ibvs(self, obs):
        cut_pose:sapien.Pose = obs['cut_pose']
        phase1_related_pose = sapien.Pose(p=[-0.1, -0.25, 0.15], q=R.from_euler('xyz', [0, 1.5708, 0.7]).as_quat()[[3, 0, 1, 2]])
        phase1_target_pose = cut_pose * phase1_related_pose

        result = self.plan_slide_pose(phase1_target_pose, obs['qpos'])
        return result
    
    def phase1_ibvs(self, obs):
        ht = obs['cut_pose'].p[2] + 0.
        slide_top = obs['slide_top']
        x_t, y_t, theta_t = 600, 480/2, 0
        x_c, y_c, theta_c = obs['cut_feature'][0], obs['cut_feature'][1], obs['cut_feature'][2]

        # 假设camera已经垂直朝下
        camera_delta_z = - 0.05 * 0.01 * (y_c - y_t)
        camera_delta_y = - 0.05 * 0.01 * (x_c - x_t)
        camera_delta_x = 0.05 * (slide_top.p[2] - ht)
        camera_delta_roll = 0.08 * (theta_c - theta_t)
        camera_delta_pose = sapien.Pose(p=[camera_delta_x, camera_delta_y, camera_delta_z], 
                                        q=R.from_euler('xyz', [camera_delta_roll, 0, 0]).as_quat()[[3, 0, 1, 2]])

        camera_pose:sapien.Pose = obs['camera_pose']
        # print(camera_pose.to_transformation_matrix()[2, 0])
        # print(obs['slide_top'])
        camera_pose_rpy = camera_pose.get_rpy()
        camera_pose.set_rpy([camera_pose_rpy[0], np.pi/2, camera_pose_rpy[2]])
        camera_target_pose = camera_pose * camera_delta_pose
        
        result = self.plan_camera_pose(camera_target_pose, obs['qpos'])
        if result['status'] != "Success":
            print('--')
        return result


if __name__ == '__main__':
    pose = sapien.Pose(p=[2, 1, 0], q=R.from_euler('xyz', [45, 0, 0], degrees=True).as_quat()[[3, 0, 1, 2]])
    related_pose = sapien.Pose(p=[1, 1, 1], q=R.from_euler('xyz', [0., 0, 0]).as_quat()[[3, 0, 1, 2]])
    new_pose = pose * related_pose
    rotate = R.from_matrix(new_pose.to_transformation_matrix()[:3, :3]).as_euler('xyz')
    print(rotate)
    print(new_pose)
    new_pose = pose * related_pose.inv()
    rotate = R.from_matrix(new_pose.to_transformation_matrix()[:3, :3]).as_euler('xyz')
    print(rotate)
    print(new_pose)